{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Extracting Results with an Agent\n",
    "\n",
    "When running a multi-agent system to solve some task, you may want to extract the result of the system once it has reached termination. This guide showcases one way to achieve this. Given that agent instances are not directly accessible from the outside, we will use an agent to publish the final result to an accessible location.\n",
    "\n",
    "If you model your system to publish some `FinalResult` type then you can create an agent whose sole job is to subscribe to this and make it available externally. For simple agents like this the {py:class}`~autogen_core.components.ClosureAgent` is an option to reduce the amount of boilerplate code. This allows you to define a function that will be associated as the agent's message handler. In this example, we're going to use a queue shared between the agent and the external code to pass the result.\n",
    "\n",
    "```{note}\n",
    "When considering how to extract results from a multi-agent system, you must always consider the subscriptions of the agent and the topics they publish to.\n",
    "This is because the agent will only receive messages from topics it is subscribed to.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "from dataclasses import dataclass\n",
    "\n",
    "from autogen_core import (\n",
    "    ClosureAgent,\n",
    "    ClosureContext,\n",
    "    DefaultSubscription,\n",
    "    DefaultTopicId,\n",
    "    MessageContext,\n",
    "    SingleThreadedAgentRuntime,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define a dataclass for the final result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class FinalResult:\n",
    "    value: str"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a queue to pass the result from the agent to the external code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "queue = asyncio.Queue[FinalResult]()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a function closure for outputting the final result to the queue.\n",
    "The function must follow the signature\n",
    "`Callable[[AgentRuntime, AgentId, T, MessageContext], Awaitable[Any]]`\n",
    "where `T` is the type of the message the agent will receive.\n",
    "You can use union types to handle multiple message types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def output_result(_agent: ClosureContext, message: FinalResult, ctx: MessageContext) -> None:\n",
    "    await queue.put(message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a runtime and register a {py:class}`~autogen_core.components.ClosureAgent` that will publish the final result to the queue."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AgentType(type='output_result')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "runtime = SingleThreadedAgentRuntime()\n",
    "await ClosureAgent.register_closure(\n",
    "    runtime, \"output_result\", output_result, subscriptions=lambda: [DefaultSubscription()]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can simulate the collection of final results by publishing them directly to the runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "runtime.start()\n",
    "await runtime.publish_message(FinalResult(\"Result 1\"), DefaultTopicId())\n",
    "await runtime.publish_message(FinalResult(\"Result 2\"), DefaultTopicId())\n",
    "await runtime.stop_when_idle()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can take a look at the queue to see the final result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Result 1\n",
      "Result 2\n"
     ]
    }
   ],
   "source": [
    "while not queue.empty():\n",
    "    print((result := await queue.get()).value)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
